{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b63c7dd3",
   "metadata": {},
   "source": [
    "# Tutorial 06, case 1b: Poisson problem with distributed control\n",
    "\n",
    "In this tutorial we solve the optimal control problem\n",
    "\n",
    "$$\\min J(y, u) = \\frac{1}{2} \\int_{\\Omega} (y - y_d)^2 dx + \\frac{\\alpha}{2} \\int_{\\Omega} u^2 dx$$\n",
    "s.t.\n",
    "$$\\begin{cases}\n",
    "- \\Delta y = f + u     & \\text{in } \\Omega\\\\\n",
    "         y = 1         & \\text{on } \\partial\\Omega\n",
    "\\end{cases}$$\n",
    "\n",
    "where\n",
    "$$\\begin{align*}\n",
    "& \\Omega               & \\text{unit square}\\\\\n",
    "& u \\in L^2(\\Omega)    & \\text{control variable}\\\\\n",
    "& y \\in H^1(\\Omega)    & \\text{state variable}\\\\\n",
    "& \\alpha > 0           & \\text{penalization parameter}\\\\\n",
    "& y_d                  & \\text{a piecewise constant desired state}\\\\\n",
    "& f                    & \\text{forcing term}\n",
    "\\end{align*}$$\n",
    "using an adjoint formulation solved by a one shot approach.\n",
    "\n",
    "The test case is from section 5.1 of\n",
    "```\n",
    "F. Negri, G. Rozza, A. Manzoni and A. Quarteroni. Reduced Basis Method for Parametrized Elliptic Optimal Control Problems. SIAM Journal on Scientific Computing, 35(5): A2316-A2340, 2013.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "693d96d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dolfinx.fem\n",
    "import dolfinx.fem.petsc\n",
    "import dolfinx.io\n",
    "import dolfinx.mesh\n",
    "import gmsh\n",
    "import mpi4py.MPI\n",
    "import numpy as np\n",
    "import petsc4py.PETSc\n",
    "import ufl\n",
    "import viskex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbd1623b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import multiphenicsx.fem\n",
    "import multiphenicsx.fem.petsc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f0aee97",
   "metadata": {},
   "source": [
    "### Geometrical parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53507a7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "L1 = 1.0\n",
    "L2 = 3.0\n",
    "H = 1.0\n",
    "mesh_size = 0.05"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26a60d5d",
   "metadata": {},
   "source": [
    "### Mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19de05f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "gmsh.initialize()\n",
    "gmsh.model.add(\"mesh\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9db590e",
   "metadata": {},
   "outputs": [],
   "source": [
    "p0 = gmsh.model.geo.addPoint(0.0, 0.0, 0.0, mesh_size)\n",
    "p1 = gmsh.model.geo.addPoint(L1, 0.0, 0.0, mesh_size)\n",
    "p2 = gmsh.model.geo.addPoint(L1 + L2, 0.0, 0.0, mesh_size)\n",
    "p3 = gmsh.model.geo.addPoint(L1 + L2, H, 0.0, mesh_size)\n",
    "p4 = gmsh.model.geo.addPoint(L1, H, 0.0, mesh_size)\n",
    "p5 = gmsh.model.geo.addPoint(0.0, H, 0.0, mesh_size)\n",
    "l0 = gmsh.model.geo.addLine(p0, p1)\n",
    "l1 = gmsh.model.geo.addLine(p1, p4)\n",
    "l2 = gmsh.model.geo.addLine(p4, p5)\n",
    "l3 = gmsh.model.geo.addLine(p5, p0)\n",
    "l4 = gmsh.model.geo.addLine(p1, p2)\n",
    "l5 = gmsh.model.geo.addLine(p2, p3)\n",
    "l6 = gmsh.model.geo.addLine(p3, p4)\n",
    "line_loop_rectangle_left = gmsh.model.geo.addCurveLoop([l0, l1, l2, l3])\n",
    "line_loop_rectangle_right = gmsh.model.geo.addCurveLoop([l4, l5, l6, -l1])\n",
    "rectangle_left = gmsh.model.geo.addPlaneSurface([line_loop_rectangle_left])\n",
    "rectangle_right = gmsh.model.geo.addPlaneSurface([line_loop_rectangle_right])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "505af6ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "gmsh.model.geo.synchronize()\n",
    "gmsh.model.addPhysicalGroup(1, [l0, l4, l5, l6, l2, l3], 1)\n",
    "gmsh.model.addPhysicalGroup(2, [rectangle_left], 1)\n",
    "gmsh.model.addPhysicalGroup(2, [rectangle_right], 2)\n",
    "gmsh.model.mesh.generate(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "741a329e",
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh, subdomains, boundaries = dolfinx.io.gmshio.model_to_mesh(\n",
    "    gmsh.model, comm=mpi4py.MPI.COMM_WORLD, rank=0, gdim=2)\n",
    "gmsh.finalize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22a1fb0b-86bb-4c3b-a746-c22b71e3b977",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create connectivities required by the rest of the code\n",
    "mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53353b6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "boundaries_1 = boundaries.indices[boundaries.values == 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "928c3518",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define associated measures\n",
    "dx = ufl.Measure(\"dx\", subdomain_data=subdomains)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff4d1d4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_mesh(mesh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4ce7816",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_mesh_tags(mesh, subdomains, \"subdomains\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7e84f6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_mesh_tags(mesh, boundaries, \"boundaries\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "531e9bfc",
   "metadata": {},
   "source": [
    "### Function spaces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c71a3768",
   "metadata": {},
   "outputs": [],
   "source": [
    "Y = dolfinx.fem.functionspace(mesh, (\"Lagrange\", 1))\n",
    "U = dolfinx.fem.functionspace(mesh, (\"Lagrange\", 1))\n",
    "Q = Y.clone()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e5127d2",
   "metadata": {},
   "source": [
    "### Trial and test functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da48f59d",
   "metadata": {},
   "outputs": [],
   "source": [
    "(y, u, p) = (ufl.TrialFunction(Y), ufl.TrialFunction(U), ufl.TrialFunction(Q))\n",
    "(z, v, q) = (ufl.TestFunction(Y), ufl.TestFunction(U), ufl.TestFunction(Q))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c525acc",
   "metadata": {},
   "source": [
    " ### Problem data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e39badda",
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha = 0.01\n",
    "y_d_1 = 1.0\n",
    "y_d_2 = 0.6\n",
    "ff = dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0))\n",
    "bc0 = petsc4py.PETSc.ScalarType(0)\n",
    "bc1 = petsc4py.PETSc.ScalarType(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7682f715",
   "metadata": {},
   "source": [
    "### Optimality conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70c19081",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = [[ufl.inner(y, z) * dx, None, ufl.inner(ufl.grad(p), ufl.grad(z)) * dx],\n",
    "     [None, alpha * ufl.inner(u, v) * dx, - ufl.inner(p, v) * dx],\n",
    "     [ufl.inner(ufl.grad(y), ufl.grad(q)) * dx, - ufl.inner(u, q) * dx, None]]\n",
    "f = [ufl.inner(y_d_1, z) * dx(1) + ufl.inner(y_d_2, z) * dx(2),\n",
    "     None,\n",
    "     ufl.inner(ff, q) * dx]\n",
    "a[2][2] = dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0)) * ufl.inner(p, q) * dx\n",
    "f[1] = ufl.inner(dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0)), v) * dx\n",
    "a_cpp = dolfinx.fem.form(a)\n",
    "f_cpp = dolfinx.fem.form(f)\n",
    "bdofs_Y_1 = dolfinx.fem.locate_dofs_topological(Y, mesh.topology.dim - 1, boundaries_1)\n",
    "bdofs_Q_1 = dolfinx.fem.locate_dofs_topological(Q, mesh.topology.dim - 1, boundaries_1)\n",
    "bc = [dolfinx.fem.dirichletbc(bc1, bdofs_Y_1, Y),\n",
    "      dolfinx.fem.dirichletbc(bc0, bdofs_Q_1, Q)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f03ac74",
   "metadata": {},
   "source": [
    "### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c762ddc",
   "metadata": {},
   "outputs": [],
   "source": [
    "(y, u, p) = (dolfinx.fem.Function(Y), dolfinx.fem.Function(U), dolfinx.fem.Function(Q))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b089efb2",
   "metadata": {},
   "source": [
    "### Cost functional"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "688c2e71",
   "metadata": {},
   "outputs": [],
   "source": [
    "J = (0.5 * ufl.inner(y - y_d_1, y - y_d_1) * dx(1) + 0.5 * ufl.inner(y - y_d_2, y - y_d_2) * dx(2)\n",
    "     + 0.5 * alpha * ufl.inner(u, u) * dx)\n",
    "J_cpp = dolfinx.fem.form(J)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca0abbeb",
   "metadata": {},
   "source": [
    "### Uncontrolled functional value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccfe2347",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract state forms from the optimality conditions\n",
    "a_state = ufl.replace(a[2][0], {q: z})\n",
    "f_state = ufl.replace(f[2], {q: z})\n",
    "a_state_cpp = dolfinx.fem.form(a_state)\n",
    "f_state_cpp = dolfinx.fem.form(f_state)\n",
    "bc_state = [bc[0]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c7f81db",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assemble the linear system for the state\n",
    "A_state = dolfinx.fem.petsc.assemble_matrix(a_state_cpp, bcs=bc_state)\n",
    "A_state.assemble()\n",
    "F_state = dolfinx.fem.petsc.assemble_vector(f_state_cpp)\n",
    "dolfinx.fem.petsc.apply_lifting(F_state, [a_state_cpp], [bc_state])\n",
    "F_state.ghostUpdate(addv=petsc4py.PETSc.InsertMode.ADD, mode=petsc4py.PETSc.ScatterMode.REVERSE)\n",
    "dolfinx.fem.petsc.set_bc(F_state, bc_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63d91e74",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve\n",
    "ksp = petsc4py.PETSc.KSP()\n",
    "ksp.create(mesh.comm)\n",
    "ksp.setOperators(A_state)\n",
    "ksp.setType(\"preonly\")\n",
    "ksp.getPC().setType(\"lu\")\n",
    "ksp.getPC().setFactorSolverType(\"mumps\")\n",
    "ksp.setFromOptions()\n",
    "ksp.solve(F_state, y.x.petsc_vec)\n",
    "y.x.petsc_vec.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)\n",
    "ksp.destroy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5027b848",
   "metadata": {},
   "outputs": [],
   "source": [
    "J_uncontrolled = mesh.comm.allreduce(dolfinx.fem.assemble_scalar(J_cpp), op=mpi4py.MPI.SUM)\n",
    "print(\"Uncontrolled J =\", J_uncontrolled)\n",
    "assert np.isclose(J_uncontrolled, 0.24)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "decaaa22",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(y, \"uncontrolled state\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e58c13ef",
   "metadata": {},
   "source": [
    "### Optimal control"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a4ce0e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assemble the block linear system for the optimality conditions\n",
    "A = dolfinx.fem.petsc.assemble_matrix_block(a_cpp, bcs=bc)\n",
    "A.assemble()\n",
    "F = dolfinx.fem.petsc.assemble_vector_block(f_cpp, a_cpp, bcs=bc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e89eee12",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve\n",
    "yup = dolfinx.fem.petsc.create_vector_block(f_cpp)\n",
    "ksp = petsc4py.PETSc.KSP()\n",
    "ksp.create(mesh.comm)\n",
    "ksp.setOperators(A)\n",
    "ksp.setType(\"preonly\")\n",
    "ksp.getPC().setType(\"lu\")\n",
    "ksp.getPC().setFactorSolverType(\"mumps\")\n",
    "ksp.setFromOptions()\n",
    "ksp.solve(F, yup)\n",
    "yup.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)\n",
    "ksp.destroy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b7cad01",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split the block solution in components\n",
    "with multiphenicsx.fem.petsc.BlockVecSubVectorWrapper(yup, [Y.dofmap, U.dofmap, Q.dofmap]) as yup_wrapper:\n",
    "    for yup_wrapper_local, component in zip(yup_wrapper, (y, u, p)):\n",
    "        with component.x.petsc_vec.localForm() as component_local:\n",
    "            component_local[:] = yup_wrapper_local"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14052b4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "J_controlled = mesh.comm.allreduce(dolfinx.fem.assemble_scalar(J_cpp), op=mpi4py.MPI.SUM)\n",
    "print(\"Optimal J =\", J_controlled)\n",
    "assert np.isclose(J_controlled, 0.158485065)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e632c56f",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(y, \"state\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4239c8a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(u, \"control\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94b2609d",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(p, \"adjoint\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython"
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
